import json
import os

import pymysql
from openai import OpenAI

client = OpenAI(api_key=os.getenv('DASHSCOPE_API_KEY'),
                base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")

tables_info = """
CREATE TABLE `emp` (
  `empno` int DEFAULT NULL, --员工编号, 默认为空
  `ename` varchar(50) DEFAULT NULL, --员工姓名, 默认为空
  `job` varchar(50) DEFAULT NULL,--员工工作, 默认为空
  `mgr` int DEFAULT NULL,--员工领导, 默认为空
  `hiredate` date DEFAULT NULL,--员工入职日期, 默认为空
  `sal` int DEFAULT NULL,--员工的月薪, 默认为空
  `comm` int DEFAULT NULL,--员工年终奖, 默认为空
  `deptno` int DEFAULT NULL,--员工部分编号, 默认为空
);

CREATE TABLE `DEPT` (
  `DEPTNO` int NOT NULL, -- 部门编码, 默认为空
  `DNAME` varchar(14) DEFAULT NULL,--部门名称, 默认为空
  `LOC` varchar(13) DEFAULT NULL,--地点, 默认为空
  PRIMARY KEY (`DEPTNO`)
);

"""
tools = [
    {
        "type": "function",
        "function": {
            "name": "ask_database",
            "description": "使用此函数回答业务问题，要求输入是mysql的查询语句",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": f"SQL查询提取信息以回答用户的问题。"
                                       f"SQL应该使用以下数据库模式编写:{tables_info}"
                                       f"查询应该以纯文本返回，而不是JSON。"
                                       f"查询应该只包含MySQL支持的语法。",
                    }
                },
                "required": ["query"],
            },
        }
    }
]


def ask_database(query):
    """连接数据库，进行查询"""
    print("进入函数内部")
    conn = pymysql.connect(
        host='localhost',
        port=3306,
        user='root',
        password='123456',
        database='edu_rag',
        charset='utf8mb4',
    )
    cursor = conn.cursor()
    cursor.execute(query)
    result = cursor.fetchall()
    cursor.close()
    conn.close()
    return result


def _execute_function_call(tool_call):
    """
    执行函数调用
    """
    available_functions = {
        "ask_database": ask_database
    }
    function_name = tool_call.function.name
    arguments = json.loads(tool_call.function.arguments)
    fn = available_functions.get(function_name)
    if fn:
        return fn(arguments.get("query"))
    else:
        return json.dumps({"error": f"Unknown function {function_name}"}, ensure_ascii=False)


def chat_completion_request(messages, tools=None, tool_choice=None, model="qwen-plus"):
    try:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            tools=tools,
            tool_choice=tool_choice,
        )
        return response
    except Exception as e:
        print("Unable to generate ChatCompletion response")
        print(f"Exception: {e}")
        return e


def _test_function_call():
    messages = [{"role": "system",
                 "content": "通过针对业务数据库生成 SQL 查询来回答用户的问题"},
                {"role": "user", "content": "查询一下最高工资的员工姓名及对应的工资"}]
    response = chat_completion_request(
        messages, tools=tools, tool_choice="auto"
    )
    assistant_message = response.choices[0].message
    print(f'assistant_message1-->{assistant_message}')

    tool_calls = response.choices[0].message.tool_calls
    if tool_calls:
        # 添加模型的响应到消息历史中
        messages.append({"role": "assistant", "content": "", "tool_calls": tool_calls})

        # 执行工具调用
        for tool_call in tool_calls:
            function_result = _execute_function_call(tool_call)
            print(f"函数调用结果:\n{function_result}")

            # 将函数结果添加到消息历史中
            messages.append({
                "role": "tool",
                "name": tool_call.function.name,
                "content": str(function_result),
                "tool_call_id": tool_call.id
            })

            # 再次调用模型，让它基于函数结果生成自然语言回复
            second_response = chat_completion_request(messages, model="qwen-plus")
            print(f"最终回复:\n{second_response.choices[0].message.content}")


if __name__ == '__main__':
    # res = ask_database('SELECT ename, sal FROM emp ORDER BY sal DESC LIMIT 1;')
    # print(res)
    _test_function_call()
    """
assistant_message1-->ChatCompletionMessage(content='', refusal=None, role='assistant', annotations=None, audio=None, function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_7f6f60552b7741c3aa5738', function=Function(arguments='{"query": "SELECT ename, sal FROM emp ORDER BY sal DESC LIMIT 1;"}', name='ask_database'), type='function', index=0)])
进入函数内部
函数调用结果:
(('David Anderson', 14955),)
最终回复:
最高工资的员工是 David Anderson，对应的工资为 14955。
    """
    pass
